# coding=utf-8

import logging

from kazoo.client import KazooClient
from kazoo.client import KazooState


class ZookeeperClient(object):
    """
    对Zookeeper操作的接口
    """

    def __init__(self, hosts, timeout=3):
        self.zk = KazooClient(hosts=hosts, timeout=timeout)
        self.zk.add_listener(zookeeper_listener)
        # 开始连接zookeeper服务
        self.zk.start()

    def createTempNode(self, path, data=b""):
        """
        建立临时路径和节点
        """
        return self.zk.create(path, data, acl=None, ephemeral=True, sequence=False, makepath=True)

    def createTempSeqNode(self, path, data=b""):
        """
        建立临时顺序节点
        """
        return self.zk.create(path, data, acl=None, ephemeral=True, sequence=True, makepath=True)

    def checkExists(self, path, watcher=None):
        """
        检查路径是否存在
        :param path:
        :param watcher:
        :return:
        """
        return self.zk.exists(path, watch=watcher) is not None

    def addWatch(self, path, watcher):
        """
        增加watch
        :param path:
        :param watcher:
        :return:
        """
        return self.zk.get(path, watch=watcher)

    def deletePath(self, path):
        """
        删除路径和节点
        """
        self.zk.delete(path)

    def getData(self, path):
        """
        获取路径数据
        """
        return self.zk.get(path)

    def setData(self, path, data):
        """
        设置节点内容
        """
        return self.zk.set(path, data)

    def getChildren(self, path):
        """
        获取指定目录下所有子节点
        """
        return self.zk.get_children(path)

    def close(self):
        """
        断开连接
        """
        self.zk.close()


def zookeeper_listener(state):
    """
    Handle being connected/reconnected to Zookeeper
    :param state:
    :return:
    """
    if state == KazooState.LOST:
        # Register somewhere that the session was lost
        logging.warn("Lost zookeeper connection!")
    elif state == KazooState.SUSPENDED:
        # Handle being disconnected from Zookeeper
        logging.warn("zookeeper connection suspended!")
    else:
        logging.warn("zookeeper connected!")
